﻿using System.Text.Json;
using Microsoft.AspNetCore.Authorization;
using Microsoft.AspNetCore.Mvc;
using WYD.Infrastructure.Controllers;
using WYD.Model.AI.OpenAIApi;
using WYD.Service.AI.IService;

namespace WYD.Admin.WebApi.Controllers.AI;

[AllowAnonymous]
[Route("v1/chat/completions")]
[ApiExplorerSettings(GroupName = "AI")]
public class OpenAIApiController : BaseController
{
    private readonly IRagService _ragService;

    public OpenAIApiController(IRagService ragService)
    {
        _ragService = ragService;
    }

    [HttpPost]
    public async Task ChatCompletions([FromBody] ChatCompletionRequest request)
    {
        // Check for an SSE stream request, which is indicated by the 'stream' field
        if (request?.Stream != true)
        {
            // If not a stream, fall back to a non-streaming response.
            // This handles clients that don't request streaming.
            await HandleNonStreamingRequest(request);
            return;
        }

        // --- Handle SSE Streaming Request ---

        // 1. Set the response headers for Server-Sent Events.
        Response.ContentType = "text/event-stream";
        Response.Headers.CacheControl = "no-cache";
        Response.Headers.Connection = "keep-alive";
        Response.Headers["X-Accel-Buffering"] = "no"; // Optional: For NGINX compatibility

        // Check for valid user input.
        // var userQuestion = request.Messages.LastOrDefault()?.Content;
        if (request.Messages.Count == 0)
        {
            // For a streaming error, send a JSON error object followed by a 'done' message.
            var errorResponse = new { error = "Request messages are missing." };
            await Response.WriteAsync($"data: {JsonSerializer.Serialize(errorResponse)}\n\n");
            await Response.WriteAsync("data: [DONE]\n\n");
            await Response.Body.FlushAsync();
            return;
        }

        try
        {
            // 2. Call the RAG service, assuming it provides an async stream of string chunks.
            // The IRagService.GetChatResponseStreamAsync method must return IAsyncEnumerable<string>.
            var createdTimestamp = DateTimeOffset.UtcNow.ToUnixTimeSeconds();
            var responseId = $"chatcmpl-{Guid.NewGuid():N}";

            await foreach (var contentChunk in _ragService.GetChatResponseStreamAsync(request.Messages))
            {
                // 3. Construct and serialize the SSE chunk data.
                var streamResponse = new ChatCompletionStreamResponse
                {
                    Id = responseId,
                    Object = "chat.completion.chunk",
                    Created = createdTimestamp,
                    Model = request.Model,
                    Choices =
                    [
                        new ChatCompletionStreamChoice
                        {
                            Index = 0,
                            Delta = new Delta { Content = contentChunk },
                            FinishReason = null
                        }
                    ],
                    // Usage = new Usage
                    // {
                    //     PromptTokens = userQuestion.Length,
                    //     CompletionTokens = ragResponse.Length,
                    //     TotalTokens = userQuestion.Length + ragResponse.Length
                    // }
                };

                var jsonString = JsonSerializer.Serialize(streamResponse, new JsonSerializerOptions
                {
                    PropertyNamingPolicy = JsonNamingPolicy.CamelCase,
                });

                // 4. Write the SSE data to the response body.
                await Response.WriteAsync($"data: {jsonString}\n\n");
                await Response.Body.FlushAsync();
            }

            // 5. Send the final [DONE] message to signal the end of the stream.
            await Response.WriteAsync("data: [DONE]\n\n");
            await Response.Body.FlushAsync();
        }
        catch (Exception ex)
        {
            // Handle exceptions within the streaming loop.
            var errorResponse = new { error = "An internal server error occurred.", details = ex.Message };
            await Response.WriteAsync($"data: {JsonSerializer.Serialize(errorResponse)}\n\n");
            await Response.WriteAsync("data: [DONE]\n\n");
            await Response.Body.FlushAsync();
        }
    }

    private async Task HandleNonStreamingRequest(ChatCompletionRequest request)
    {
        // This is the original non-streaming logic from your provided code.
        if (request.Messages.Count == 0)
        {
            Response.StatusCode = StatusCodes.Status400BadRequest;
            await Response.WriteAsJsonAsync(new { error = "Invalid request." });
            return;
        }
        
        try
        {
            var ragResponse = await _ragService.GetChatResponseAsync(request.Messages);

            var response = new ChatCompletionResponse
            {
                Id = $"chatcmpl-{Guid.NewGuid().ToString().Substring(0, 12)}",
                Object = "chat.completion",
                Created = DateTimeOffset.UtcNow.ToUnixTimeSeconds(),
                Model = request.Model,
                Choices = new List<Choice>
                {
                    new()
                    {
                        Index = 0,
                        Message = new Message
                        {
                            Role = "assistant",
                            Content = ragResponse
                        },
                        FinishReason = "stop"
                    }
                },
                // Usage = new Usage
                // {
                //     PromptTokens = userQuestion.Length,
                //     CompletionTokens = ragResponse.Length,
                //     TotalTokens = userQuestion.Length + ragResponse.Length
                // }
            };
            await Response.WriteAsJsonAsync(response);
        }
        catch (Exception ex)
        {
            Response.StatusCode = StatusCodes.Status500InternalServerError;
            await Response.WriteAsJsonAsync(new { error = "An internal server error occurred.", details = ex.Message });
        }
    }
}